import glob
import os.path

import numpy as np
import torch
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image

import cv2

# http://www.cs.toronto.edu/~kriz/cifar.html 数据集下载
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


# 'airplane'=0
# 'automobile'=1
# 'brid'=2
# 'cat'=3
# 'deer'=4
# 'dog'=5
# 'frog'=6
# 'horse'=7
# 'ship'=8
# 'truck'=9

label_name = ['airplane',
              'automobile',
              'brid',
              'cat',
              'deer',
              'dog',
              'frog',
              'horse',
              'ship',
              'truck',
              ]
label_dict = {}
for idx, name in enumerate(label_name):
    label_dict[name] = idx


def prepare_train():
    save_path = "../data/train"
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    train_list = glob.glob("../data/cifar-10-batches-py/data_batch_*")
    # print(train_list)
    for l in train_list:
        l_dict = unpickle(l)
        # print(l_dict)
        # print((l_dict.keys()))
        for im_idx, im_data in enumerate(l_dict[b'data']):
            # print(im_idx)
            # print(im_data)
            im_label = l_dict[b'labels'][im_idx]
            im_name = l_dict[b'filenames'][im_idx]
            # print(im_label)
            # print(im_name)

            im_label_name = label_name[im_label]
            im_data = np.reshape(im_data, [3, 32, 32])
            im_data = np.transpose(im_data, (1, 2, 0))

            # cv2.imshow("im_data",cv2.resize(im_data,(200,200)))
            # cv2.waitKey(0)

            im_path = "{}/{}".format(save_path, im_label_name)
            if not os.path.exists(im_path):
                os.mkdir(im_path)
            cv2.imwrite("{}/{}/{}".format(save_path, im_label_name, im_name.decode("utf-8")), im_data)


def prepare_test():
    save_path = "../data/test"
    if not os.path.exists(save_path):
        os.mkdir(save_path)

    test_list = glob.glob("../data/cifar-10-batches-py/test_batch*")
    # print(train_list)
    for l in test_list:
        l_dict = unpickle(l)
        # print(l_dict)
        # print((l_dict.keys()))
        for im_idx, im_data in enumerate(l_dict[b'data']):
            # print(im_idx)
            # print(im_data)
            im_label = l_dict[b'labels'][im_idx]
            im_name = l_dict[b'filenames'][im_idx]
            # print(im_label)
            # print(im_name)

            im_label_name = label_name[im_label]
            im_data = np.reshape(im_data, [3, 32, 32])
            im_data = np.transpose(im_data, (1, 2, 0))

            # cv2.imshow("im_data",cv2.resize(im_data,(200,200)))
            # cv2.waitKey(0)

            im_path = "{}/{}".format(save_path, im_label_name)
            if not os.path.exists(im_path):
                os.mkdir(im_path)
            cv2.imwrite("{}/{}/{}".format(save_path, im_label_name, im_name.decode("utf-8")), im_data)


def default_loader(path):
    return Image.open(path).convert("RGB")


# 对数据做一系列的增强，比如变成28*28，颜色变成灰色
# train_transform = transforms.Compose([
#     transforms.RandomResizedCrop((28, 28)),
#     transforms.RandomVerticalFlip(),
#     transforms.RandomHorizontalFlip(),
#     transforms.RandomRotation(90),
#     transforms.RandomGrayscale(0.1),
#     transforms.ColorJitter(0.3, 0.3, 0.3, 0.3),
#     transforms.ToTensor(),
# ])
train_transform = transforms.Compose([
    transforms.RandomCrop(28),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
])
test_transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
])


class MyDataset(Dataset):
    def __init__(self, im_list, transform=None, loader=default_loader):
        super(MyDataset, self).__init__()
        self.imgs = []
        for im_item in im_list:
            im_label_name = im_item.split("/")[-2]
            self.imgs.append([im_item, label_dict[im_label_name]])
        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        im_path, im_label = self.imgs[index]
        im_data = self.loader(im_path)
        if self.transform is not None:
            im_data = self.transform(im_data)
        return im_data, im_label

    def __len__(self):
        return len(self.imgs)


if __name__ == '__main__':
    prepare_train()
    prepare_test()
